import os

import numpy as np
import cv2
import torch
import torchvision
from ultralytics import YOLO

if __name__ == '__main__':
    # 自定义新模型，迁移学习，虽然与训练模型结构关键点个数与自定义数据集关键点个数不同，但是模型可以根据数据配置文件进行自适应
    model = YOLO('yolov8n-pose.pt')  # 所以最好之间迁移学习即可，模型已经在GPU，可以打断点调试下看看设备
    # 开始模型训练，注意配置好数据的yaml文件
    model.train(data=r'./plate-data.yaml', epochs=75, imgsz=640, workers=0, batch=32)

    # 对于存在的问题：
    '''
    1. 在自定义数据集上训练yolov8模型，需要准备几个文件？是模型结构文件（权重文件）和数据配置文件吗？
    答：需要准备两个文件，模型权重文件和数据配置文件，不过具体的数据需要自己预处理好。对于其中的模型权重文件，可以使用与训练模型进行迁移学习，模型会根据数据配置文件里面的yaml文件的关键点个数进行拟合，模型结构不用动
    
    2. 训练自定义数据集上的yolov8模型的时候，如果要使用迁移学习，从官方与训练模型加载模型和权重，由于自定义数据集关键点个数和格式与官方与训练数据集所使用的数据格式不同，模型结构不匹配会报错？
    答：可以正常进行训练，不过要在数据配置文件yaml文件下配置好关键点的数量
    
    3. 如果（2）会报错，是不是只需要指定.yaml文件即模型结构文件，修改输出头结构以匹配自定义数据集即可
    答：也可以使用yaml文件进行从0开始学习训练一个模型
    
    4. 但是yolov8官方只给了通用的.yaml模型结构，如何指定特定版本的n,s,l,x结构模型？
    答：把模型配置yaml文件里面的前面的版本数据留下，做成一个自己的模型配置文件结构
    
    5. 就算指定了特定结构模型，为什么会自动下载模型权重？？？比如yolov11n.pt
    答：yolov8下载的权重文件比如yolov11n.pt似乎不是用于训练的，而是一个跟训练过程中的 AMP check 有关，所以不用管，详见：https://github.com/ultralytics/ultralytics/issues/18220
    
    6. 我用的模型是yolov8，就算自动下载模型权重不是也应该下载yolov8的嘛？为什么会自动下载yolov11版本的模型呢？
    答：见5
    '''

    # 这是训练后的结果
    '''
    training result:
    75 epochs completed in 1.912 hours.
    Optimizer stripped from runs\pose\train7\weights\last.pt, 6.4MB
    Optimizer stripped from runs\pose\train7\weights\best.pt, 6.4MB
    
    Validating runs\pose\train7\weights\best.pt...
    Ultralytics 8.3.109 🚀 Python-3.12.9 torch-2.6.0+cu126 CUDA:0 (NVIDIA GeForce RTX 4060 Laptop GPU, 8188MiB)
    YOLOv8n-pose summary (fused): 81 layers, 3,077,771 parameters, 0 gradients, 8.3 GFLOPs
                     Class     Images  Instances      Box(P          R      mAP50  mAP50-95)     Pose(P          R      mAP50  mAP50-95): 100%|██████████| 16/16 [00:08<00:00,  1.82it/s]
                       all       1001       1001      0.999      0.999      0.994      0.919          1          1      0.995      0.995
    Speed: 0.1ms preprocess, 1.0ms inference, 0.0ms loss, 0.6ms postprocess per image
    Results saved to runs\pose\train7

    '''